package com.allinpay.mars.common.db.vault.config;

import com.allinpay.mars.common.db.vault.util.RestTemplateUtils;
import com.zaxxer.hikari.HikariDataSource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.ResponseEntity;
import org.springframework.scheduling.annotation.Scheduled;

import javax.sql.DataSource;
import java.net.URISyntaxException;
import java.util.HashMap;
import java.util.Map;

@Configuration
@Slf4j
public class MarsHiKariConfig {

    @Value("${spring.datasource.url}")
    private String url;

    @Autowired
    private MarsVaultConfig marsVaultConfig;

    private CustomHiKariDatasource customHiKariDatasource = new CustomHiKariDatasource();

    @Bean(name = "dataSource")
    public DataSource dataSource() throws URISyntaxException {
        log.info("init datasource");
        initDatasource();
        return customHiKariDatasource;
    }

    private void initDatasource() {
        HikariDataSource dataSource = new HikariDataSource();
        dataSource.setJdbcUrl(url);
        String cacheUsername = "";
        String cachePassword = "";
        Map<String, Map<String, Object>> cacheTokenMap = customHiKariDatasource.getTokenMap();
        boolean useCacheToken = true;
        if (cacheTokenMap.isEmpty()) {
            useCacheToken = false;
        } else {
            Map<String, Object> token = cacheTokenMap.get("token");
            Long expectExpire = (Long) token.get("expectExpire");
            Long curr = System.currentTimeMillis() / 1000;
            if (curr > expectExpire) {
                useCacheToken = false;
            } else {
                cacheUsername = (String) token.get("username");
                cachePassword = (String) token.get("password");
            }
        }
        if (useCacheToken) {
            dataSource.setUsername(cacheUsername);
            dataSource.setPassword(cachePassword);
        } else {
            getTokenFromVault();
            String username = customHiKariDatasource.getUsername();
            String password = customHiKariDatasource.getPassword();
            dataSource.setUsername(username);
            dataSource.setPassword(password);
        }
        dataSource.setPoolName("current");
        customHiKariDatasource.updateDataSourceMap("1", dataSource);
    }

    private void getTokenFromVault() {
        String vaultLink = marsVaultConfig.getUrl() + marsVaultConfig.getRole();
        HttpHeaders httpHeaders = new HttpHeaders();
        httpHeaders.set("X-Vault-Token", marsVaultConfig.getToken());
        HttpEntity<Map> entity = new HttpEntity<>(httpHeaders);
        ResponseEntity<Map> responseEntity = RestTemplateUtils.exchange(vaultLink, HttpMethod.GET, entity, Map.class);
        Map body = responseEntity.getBody();
        if (body.containsKey("data")) {
            Map data = (Map) body.get("data");
            String username = (String) data.get("username");
            String password = (String) data.get("password");
            Map<String, Object> token = new HashMap<>();
            token.put("username", username);
            token.put("password", password);
            token.put("ttl", (Integer) data.get("ttl"));
            token.put("expectExpire", System.currentTimeMillis() / 1000 + (Integer) data.get("rotation_period"));
            token.put("curr", System.currentTimeMillis() / 1000);
            customHiKariDatasource.updateTokenMap("token", token);
        }
    }

    @Scheduled(cron = "* 30 14 * * ?")
    public void update() {
        log.info("schedule refresh datasource");
        initDatasource();
    }

}
